
import numpy as np
from torch.utils.data import Dataset

class Mydata(Dataset):
    #744个od  前528个（22天）作为训练，后9天（216）做测试
    def __init__(self,taxi_od,is_train_data = True):
        super().__init__()
        self.x = []
        self.y = []
        if is_train_data:
            for i in range((22-3)*24):
                temx = np.random.rand(1,6,69,69)
                temy = np.random.rand(1,69,69)
                for j in range(3):
                    temx[0,j,:,:] = taxi_od[(72+i)-(3-j),:,:]
                for j in range(3):
                    temx[0,j+3,:,:] = taxi_od[(72+i)-(3-j)*24,:,:]
                temy[0,:,:] = taxi_od[(72+i),:,:]
                self.x.append(temx)
                self.y.append(temy)
        else:
            for i in range((9)*24):
                temx = np.random.rand(1,6,69,69)
                temy = np.random.rand(1,69,69)
                for j in range(3):
                    temx[0,j,:,:] = taxi_od[(528+i)-(3-j),:,:]
                for j in range(3):
                    temx[0,j+3,:,:] = taxi_od[(528+i)-(3-j)*24,:,:]
                temy[0,:,:] = taxi_od[(528+i),:,:]
                self.x.append(temx)
                self.y.append(temy)
    def __getitem__(self,index):
        return self.x[index],self.y[index]
    def __len__(self):
        return len(self.x)